"""Utility for measuring container height-volume calibration data on a Hamilton STAR."""

import statistics
import warnings
from typing import TYPE_CHECKING, Awaitable, Callable, Dict, List, Optional

if TYPE_CHECKING:
    from pylabrobot.liquid_handling.liquid_handler import LiquidHandler
    from pylabrobot.resources.container import Container


async def _default_prompt(volume: float, step_idx: int, total: int) -> None:
    input(f"[{step_idx + 1}/{total}] Pipette {volume} uL into the container, then press Enter...")


async def measure_height_volume_data(
    lh: "LiquidHandler",
    container: "Container",
    test_volumes: List[float],
    channel_idx: int = 0,
    # ztouch params
    ztouch_num_replicates: int = 5,
    ztouch_speed: float = 18.0,
    ztouch_detection_limiter: int = 0,
    ztouch_push_down_force: int = 0,
    # cLLD params
    clld_num_replicates: int = 3,
    clld_speed: float = 6.0,
    clld_detection_edge: int = 5,
    clld_post_detection_dist: float = 2.0,
    # output
    decimal_places: int = 2,
    prompt_callback: Optional[Callable[[float, int, int], Awaitable[None]]] = None,
) -> Dict[float, float]:
    """Measure height-volume calibration data for a container using a Hamilton STAR.

    Uses ztouch to find the cavity bottom, then cLLD at each test volume to build a
    ``{height_mm: volume_uL}`` mapping suitable for ``Container.height_volume_data``.

    The caller is responsible for tip pickup/drop. A teaching needle (metal tip) must
    already be mounted on ``channel_idx``.

    Args:
        lh: A set-up LiquidHandler with a STARBackend.
        container: The container to calibrate (must be on the deck).
        test_volumes: Sorted list of positive volumes (uL) to measure.
        channel_idx: Channel to use (0-based).
        ztouch_num_replicates: Number of ztouch replicates for cavity bottom.
        ztouch_speed: Ztouch probe speed in mm/s.
        ztouch_detection_limiter: PWM limiter for ztouch detection.
        ztouch_push_down_force: PWM push-down force for ztouch.
        clld_num_replicates: Number of cLLD replicates per volume.
        clld_speed: cLLD probe speed in mm/s.
        clld_detection_edge: Edge steepness for cLLD detection.
        clld_post_detection_dist: Distance to move after cLLD detection in mm.
        decimal_places: Rounding precision for height values.
        prompt_callback: Async callable ``(volume, step_idx, total) -> None`` invoked before
            each volume measurement. Defaults to ``input()`` prompt.

    Returns:
        Dict mapping height (mm above cavity bottom) to volume (uL), always including
        ``{0.0: 0.0}`` as the baseline.

    Raises:
        TypeError: If the backend is not a STARBackend.
        RuntimeError: If no tip is present on the channel.
        ValueError: If test_volumes is invalid (empty, unsorted, duplicates, or negatives).
    """

    from pylabrobot.liquid_handling.backends.hamilton.STAR_backend import STARBackend

    # --- Validation (no hardware touched) ---
    backend = lh.backend
    if not isinstance(backend, STARBackend):
        raise TypeError(
            f"measure_height_volume_data requires a STARBackend, got {type(backend).__name__}"
        )

    if not test_volumes:
        raise ValueError("test_volumes must not be empty")
    if any(v <= 0 for v in test_volumes):
        raise ValueError("All test_volumes must be positive")
    if test_volumes != sorted(test_volumes):
        raise ValueError("test_volumes must be sorted in ascending order")
    if len(test_volumes) != len(set(test_volumes)):
        raise ValueError("test_volumes must not contain duplicates")

    if not (0 <= channel_idx < backend.num_channels):
        raise ValueError(f"channel_idx {channel_idx} out of range [0, {backend.num_channels})")

    if not lh.head[channel_idx].has_tip:
        raise RuntimeError(f"No tip present on channel {channel_idx}")

    if prompt_callback is None:
        prompt_callback = _default_prompt

    # --- Position channel over container ---
    await backend.move_all_channels_in_z_safety()
    await backend.prepare_for_manual_channel_operation(channel_idx)

    container_center = container.get_absolute_location("c", "c", "t")
    await backend.move_channel_x(channel_idx, container_center.x)
    await backend.move_channel_y(channel_idx, container_center.y)

    # --- Measurement with safety wrapper ---
    result: Dict[float, float] = {0.0: 0.0}

    try:
        # Phase A: Cavity bottom (ztouch)
        ztouch_readings = []
        last_z: Optional[float] = None

        for i in range(ztouch_num_replicates):
            start_pos = None
            if last_z is not None:
                start_pos = last_z + 5.0

            z = await backend.ztouch_probe_z_height_using_channel(
                channel_idx=channel_idx,
                channel_speed=ztouch_speed,
                detection_limiter_in_PWM=ztouch_detection_limiter,
                push_down_force_in_PWM=ztouch_push_down_force,
                start_pos_search=start_pos,
                move_channels_to_safe_pos_after=True,
            )
            ztouch_readings.append(z)
            last_z = z

        cavity_bottom_z = statistics.mean(ztouch_readings)
        if len(ztouch_readings) > 1:
            std = statistics.stdev(ztouch_readings)
            if std > 0.5:
                warnings.warn(
                    f"High ztouch variability: std={std:.3f}mm across "
                    f"{ztouch_num_replicates} replicates. Readings: {ztouch_readings}",
                    stacklevel=2,
                )

        print(
            f"Cavity bottom Z: {cavity_bottom_z:.{decimal_places}f} mm "
            f"(std={statistics.stdev(ztouch_readings) if len(ztouch_readings) > 1 else 0:.3f}mm)"
        )

        # Phase B: cLLD measurements per volume
        last_clld_z: Optional[float] = None
        prev_height: Optional[float] = None

        for step_idx, volume in enumerate(test_volumes):
            await prompt_callback(volume, step_idx, len(test_volumes))

            clld_readings = []
            for j in range(clld_num_replicates):
                start_pos = None
                if last_clld_z is not None:
                    start_pos = last_clld_z + 5.0

                z = await backend.clld_probe_z_height_using_channel(
                    channel_idx=channel_idx,
                    channel_speed=clld_speed,
                    detection_edge=clld_detection_edge,
                    post_detection_dist=clld_post_detection_dist,
                    start_pos_search=start_pos,
                    move_channels_to_safe_pos_after=True,
                )
                clld_readings.append(z)
                last_clld_z = z

            mean_z = statistics.mean(clld_readings)
            height = round(mean_z - cavity_bottom_z, decimal_places)

            if prev_height is not None and height <= prev_height:
                warnings.warn(
                    f"Non-monotonic height at {volume} uL: {height}mm <= previous {prev_height}mm. "
                    f"Check for bubbles or evaporation.",
                    stacklevel=2,
                )

            result[height] = volume
            prev_height = height

            std_str = (
                f", std={statistics.stdev(clld_readings):.3f}mm"
                if len(clld_readings) > 1
                else ""
            )
            print(f"  {volume:>8.1f} uL -> {height:.{decimal_places}f} mm{std_str}")

    finally:
        try:
            await backend.move_all_channels_in_z_safety()
        except Exception:
            pass

    # Phase C: Output
    result = dict(sorted(result.items(), key=lambda item: item[0]))

    print("\n# Copy-paste into your container definition:")
    print("height_volume_data = {")
    for h, v in result.items():
        print(f"  {h}: {v},")
    print("}")

    min_h = min(result.keys())
    max_h = max(result.keys())
    min_v = min(result.values())
    max_v = max(result.values())
    print(
        f"\nSummary: {len(result)} points, "
        f"height [{min_h:.{decimal_places}f}, {max_h:.{decimal_places}f}] mm, "
        f"volume [{min_v:.1f}, {max_v:.1f}] uL"
    )

    return result